# -*- coding: UTF-8 -*-
# @File : serve.py
# @IDE     : VScode
# @Author  : zhonggc
# @Date    : 2024/05/24 16:17:26

from typing import List
from fastapi import FastAPI
from langchain_core.prompts import ChatPromptTemplate
from langchain_core.output_parsers import StrOutputParser
from langchain_cohere import ChatCohere
from langserve import add_routes
from dotenv import load_dotenv, find_dotenv

# 1.Create prompt template
system_template = "Translate the following into {language}:"
prompt_template = ChatPromptTemplate.from_messages(
    [("system", system_template), ("user", "{text}")]
)

# 2.Create model
model = ChatCohere(model="command-r")

# 3.Create parser
parser = StrOutputParser()

# 4.Create chain
chain = prompt_template | model | parser

# 5.Create app
app = FastAPI(
    title="langChain Sserver",
    version="1.0",
    description="A simple langChain API server",
)

# 6.Adding chain routes
add_routes(
    app,
    chain,
    path="/chain",
)


if __name__ == "__main__":
    import uvicorn

    _ = load_dotenv(find_dotenv())
    uvicorn.run(app, host="localhost", port=8000)
